import SparkApi
import re
import json

#用于配置大模型版本，默认“general/generalv2”
domain = "generalv3"    # v3.0版本
#云端环境的服务地址
Spark_url = "ws://spark-api.xf-yun.com/v3.1/chat"  # v3.0环境的地址


def process_llm_history(config, paras): # under development, if we need to parse the history
    his_record = ''
    try:
        his_k = config['HISTORY_K']
        histories = paras['histories']
        
        if len(histories)>=2:
            his_infos = histories[-2 :]
            for i, his in enumerate(his_infos):
                his_record += f"""历史提问{i+1}：{his[0]}    历史回答{i+1}：{his[1]}\n"""
            his_record = his_record.strip()
    except:
        pass
    
    return his_record

       
def getText(role, content):
    jsoncon = {}
    jsoncon["role"] = role
    jsoncon["content"] = content
    return jsoncon


def eval_response(str_answer, config=None, answer_key=None):
    str_answer = re.sub(r'\s+', '', str_answer)        
    str_answer = str_answer.replace("'", '"').strip()
    print('\t' + str_answer + '\n')
    
    try:
        answer = json.loads(str_answer)
        print('\t successfully find json marks and parse the response')
    except:
        print('\t regular parsing fails, try finding json marks...')

        matches = re.findall(r'```json\s*(.+?)\s*```', str_answer, re.DOTALL)
        if len(matches)==0:
            matches = re.findall(r'```json({.+?})```', str_answer, re.DOTALL)
            if len(matches)==0:
                str_answer = str_answer.replace('```json', '').strip()
                str_answer = str_answer.replace('```', '').strip()
            else:
                str_answer = matches[0]
        else:
            str_answer = matches[0]
            
        try:
            answer = json.loads(str_answer)
        except:
            print('\t no json object parsed, using original input answer')
            return str_answer
         
    if not answer_key==None:
        answer = answer[answer_key]
    return answer


def call_spark(**kwargs):
    try:
        appid = kwargs['config']['SPARK_APPID']
        api_secret = kwargs['config']['SPARK_SECRET']
        api_key =kwargs['config']['SPARK_APIKEY']
    except:
        print('Bad config inputs...')
        return None
    
    try:
        task = kwargs['paras']['task']
        query = kwargs['paras']['query']
        texts = kwargs['paras']['texts']
    except:
        print('Bad paras inputs...')
        return None
    
    his_record = process_llm_history(kwargs['config'], kwargs['paras'])
            
    if task=='talk':
        add_req = kwargs['paras']['add_req']
        out_limit = kwargs['paras']['out_limit']
        
        prompt = f"""
                你将接收到一些使用json字典格式存储的基础资料：'''{texts}'''。
                你接收到用户的问题：'''{query}'''。
                你可能还会接收到历史对话中，你曾给用户提供的相关答案：'''{his_record}'''。
                请你根据已有资料和历史记录（如有），回答用户问题。{add_req}，不要超过{out_limit}字。
            """

    elif task=='talk-kb':
        add_req = kwargs['paras']['add_req']
        if_judge = kwargs['paras']['if_judge']
        out_limit = kwargs['paras']['out_limit']
        
        if if_judge:
            example = '''
                    {
                        "judge": "否",
                        "answer": "根据您的资料答案是..."
                    }
                '''
            
            prompt = f"""
                    你将接收到用户的问题：'''{query}'''
                    你将接收到一些基础知识和资料：'''{texts}'''
                    你可能还会接收到历史对话中的相关问题和答案：'''{his_record}'''
                    你要执行以下两步：
                    1）判断输入的基础知识资料中是否包含了用户问题的答案。
                    2）根据基础知识资料以及历史记录（如有）回答用户问题。{add_req}，不要超过{out_limit}字。
                    
                    你的返回必须是JSON字典格式，包含2个键，第1个键是"judge"，值为输入的基础知识资料中是否包含了用户问题的答案，只能“是”或“否"；第2个键是"answer"，值为你的回答，下面是一个返回示例
                """    
            prompt = prompt + example
        else:
            prompt = f"""
                    你将接收到用户的问题：'''{query}'''
                    你将接收到一些基础知识和资料：'''{texts}'''
                    你可能还会接收到历史对话中的相关问题和答案：'''{his_record}'''
                    请严格参考基础知识资料以及历史记录（如有），回答用户问题。{add_req}，不要超过{out_limit}字。
                """
            
    elif task=='judge': # Patent idea: feature based prompt and agent management
        structure = kwargs['structure']
        prompt = f"""
                假设你是一名在科研机构工作者，你的任务是对接企业和科研机构的研究成果。
                你现在将接收到一个描述团队及其研究成果的json树形结构：'''{structure}'''。可见，其中包含了团队、相关技术文件名字、技术文件中的重要标题名字等。
                你还接收到用户的问题或需求：'''{query}'''。
                
                你需要理解该树形结构中的键，判断使用哪些键最能够回答用户的问题，对基于此修剪输入的树形结构。
                比如，用户提问中涉及到具体团队或者文件名称或次序，你应该保留对应的特定团队、文件名、标题等，并剔除其他非相关项；又比如，用户提问是全局性的，例如进行总体概括等，则你可以保留全部的键
                
                你的最终回答只能是json格式，你的回答的主键是"answer"，值为根据你的判断进行过修剪的json结构.
                除了json格式的返回之外，不要附带其他信息或解释。
            """
    elif task=='summary':
        try:
            limit = kwargs['limit']
            add_req = kwargs['add_req']
        except:
            limit = int(len(texts)*0.2)
            add_req = ''
        prompt = f"""
                你将接收一段文字材料：{texts} {add_req}
                你需要总结以上文字材料主要内容，字数不超过{limit}字。
                你的回答只包含提炼后的文本，不要返回其他多余信息。
            """
    elif task=='rewrite':
        style = kwargs['style']           
        prompt = f"""
                对下面这段输入的文本内容进行重写，使其符合{style}的风格：'''{texts}'''
                你的回答只能是json形式，键为"answer"，值为重写后的文本，不要返回其他多余信息。
            """
    elif task=='repair json':         
        prompt = f"""
                你会接收一段json字典形式的文本：{texts}.
                请你判断这个json的语法格式是否正确，如有错误，请进行修改并返回修改后的json字典。
                如无错误，直接返回原始json字典。
                无论如何，你的回答只能是json形式，不要返回其他多余信息。
            """
    
    elif task=='inspect-typos':
        example = '''
                如果用户的输入存在错别字，例如“令天我们去工地按装一台设备”，则你的回答示例如下：
                {
                    "回答" : [
                        {"原词": "令天", "修正":"今天" },
                        {"原词": "按装", "修正":"安装" }
                    ]
                }
                
                如果用户输入不存在错别字，则修改列表为空，例如“工程项目的付款周期为3个月” 你的回答示例如下：
                {
                    "回答" : [ ]
                }
            '''
        
        prompt = f"""
                你会收到一段关于文字，你需要检查文中错别字，但是对可能的专有名词，默认正确，不要修改。
                你的回答必须是JSON形式，键为“回答”，值为一个列表，列表中每个元素均为字典，包含2个键，即"原词"和"修正"，对应的值为存在错别字的短语以及修改后的结果。
                注意，除以上形式的JSON数据外，不要返回其他任何内容或解释。返回示例如下：
                
                {example}
                
                现在，按以上规则和示例，分析下文错别字： '''{texts}'''
            """   
            
    else:
        pass
    
    # print(prompt)
    SparkApi.answer = '' # reset
    spark_query = getText("user", prompt)
    
    SparkApi.ask_spark(appid, api_key, api_secret, Spark_url, domain, [spark_query])
    final_answer = eval_response(SparkApi.answer)
    
    return final_answer
    
            
    
    


